package com.pitaya.avatar.util;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import com.pitaya.avatar.service.IClusterNode;
import com.pitaya.avatar.service.impl.Distance;
import com.pitaya.avatar.service.impl.Node;
import com.pitaya.avatar.service.impl.UnionFindSet;

public class KruskalCluster {
	
	private UnionFindSet ufSet;
	
	public void setUfSet(UnionFindSet ufset){
		this.ufSet = ufset;
	}
	
	public UnionFindSet getUfSet(){
		return this.ufSet;
	}
	
	public ArrayList<ArrayList<IClusterNode>> clusterNodes(ArrayList<? extends Distance> distances, ArrayList<Node> nodes, int k){
		
		 Map<String, ArrayList<IClusterNode>> nodesClusters = new HashMap<String, ArrayList<IClusterNode>>();
		 ArrayList<ArrayList<IClusterNode>> resultList = new ArrayList<ArrayList<IClusterNode>>();
		 Collections.sort(distances);
		 int times = nodes.size() - 1 - (k - 1); 
		 int count = 0;
		 for(Distance distance : distances){
			 if(count == times)
				 break;
			 Node root1 = ufSet.find(distance.getStartNode());
			 Node root2 = ufSet.find(distance.getEndNode());
			 if(root1.equals(root2)){
				 continue;
			 } else {
				 count ++;
				 ufSet.union(root1, root2);
			 }
		 }
		 //Generate the result list
		 for(Node node : nodes){
			 String setName = ufSet.find(node).getNodeName();
			 ArrayList<IClusterNode> nodesCluster = nodesClusters.get(setName);
			 if(nodesCluster == null){
				 nodesCluster = new ArrayList<IClusterNode>();
				 nodesClusters.put(setName, nodesCluster);
			 }
			 nodesCluster.add(node.getClusterNode());
		 }
		 for(Map.Entry<String, ArrayList<IClusterNode>> entry : nodesClusters.entrySet()){
			 resultList.add(entry.getValue());
		 }
		 return resultList;
	}
	
}
